#!/usr/bin/env python3
"""
Script to split statcomp.tex into 9 separate files:
- Files 1-8: sections 3.1 to 3.8
- File 9: remaining content (习题三 and other sections)
- main.tex: main file that includes all sections
"""

import re

def split_tex_file(input_file, output_prefix="section"):
    """Split LaTeX file by sections."""

    # Read the entire file
    with open(input_file, 'r', encoding='utf-8') as f:
        content = f.read()

    # Find the preamble (everything before section 3.1)
    match = re.search(r'\\section\*\{第三章', content)
    if not match:
        print("Error: Could not find chapter section")
        return

    preamble = content[:match.start()]
    chapter_title_match = re.search(r'\\section\*\{第三章 随机模拟\}', content)
    chapter_title = content[chapter_title_match.start():chapter_title_match.end()]

    # Define section markers for 3.1 through 3.8
    section_patterns = [
        (r'\\section\*\{3\.1', '3.1'),
        (r'\\section\*\{3\.2 ', '3.2'),
        (r'\\section\*\{3\.3 ', '3.3'),
        (r'\\section\*\{3\.4 ', '3.4'),
        (r'\\section\*\{3\.5 ', '3.5'),
        (r'\\section\*\{3\.6 ', '3.6'),
        (r'\\section\*\{3\.7 ', '3.7'),
        (r'\\section\*\{3\.8 ', '3.8'),
    ]

    # Find all section positions
    sections = []
    for pattern, section_num in section_patterns:
        match = re.search(pattern, content)
        if match:
            sections.append((match.start(), section_num))
        else:
            print(f"Warning: Could not find section {section_num}")

    # Sort by position
    sections.sort()

    # Find where exercises start (习题三)
    exercises_match = re.search(r'\\section\*\{习题三\}', content)
    if exercises_match:
        exercises_start = exercises_match.start()
    else:
        print("Warning: Could not find 习题三")
        exercises_start = len(content)

    # List to store section file names
    section_files = []

    # Extract and write each section file
    for i, (start_pos, section_num) in enumerate(sections):
        # Determine end position
        if i < len(sections) - 1:
            end_pos = sections[i + 1][0]
        else:
            # Last section (3.8) ends at exercises
            end_pos = exercises_start

        # Extract section content (without preamble)
        section_content = content[start_pos:end_pos].rstrip() + "\n"

        # Write to file
        output_file = f"{output_prefix}_{section_num}.tex"
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(section_content)

        section_files.append(output_file)
        print(f"Created {output_file}")

    # Create the 9th file with remaining content (exercises and other sections)
    end_doc_pos = content.find(r'\end{document}')
    remaining_content = content[exercises_start:end_doc_pos].rstrip() + "\n"

    output_file = f"{output_prefix}_9_exercises.tex"
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(remaining_content)

    section_files.append(output_file)
    print(f"Created {output_file}")

    # Create main.tex file
    main_content = preamble + chapter_title + "\n\n"
    for section_file in section_files:
        main_content += f"\\input{{{section_file}}}\n"
    main_content += "\n\\end{document}\n"

    with open("main.tex", 'w', encoding='utf-8') as f:
        f.write(main_content)

    print("Created main.tex")
    print(f"\nSplit complete! Created 9 section files + 1 main file.")

if __name__ == "__main__":
    split_tex_file("statcomp.tex")
